import numpy as np 
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# 洛伦兹系统方程
def lorenz(x, y, z, s=10, r=28, b=2.667):
    x_dot = s*(y - x)
    y_dot = r*x - y - x*z 
    z_dot = x*y - b*z
    return x_dot, y_dot, z_dot

# 运行洛伦兹系统,返回轨迹
def run_lorenz(x0, y0, z0, t, dt):
    x = np.empty((t,))
    y = np.empty((t,))
    z = np.empty((t,))
    x[0] = x0
    y[0] = y0
    z[0] = z0
    
    for i in range(t-1):
        x_dot, y_dot, z_dot = lorenz(x[i], y[i], z[i])
        x[i+1] = x[i] + x_dot*dt
        y[i+1] = y[i] + y_dot*dt 
        z[i+1] = z[i] + z_dot*dt
        
    return x, y, z 

def calculate_lyapunov(x, y, z, dt):
    beta = 8/3  # 定义beta为常数
    n = len(x)
    m = np.identity(3)
    v = np.zeros(3)
    d = 0
    for i in range(n-1):
        x_dot, y_dot, z_dot = lorenz(x[i], y[i], z[i])
        j = np.array([[x_dot, y_dot, z_dot],
                      [0, 0, -x_dot],
                      [0, x_dot, -beta]])
        m = np.matmul(j, m)
        v += np.log(np.abs(np.diag(m)))
        m, r = np.linalg.qr(m)
        if i % 100 == 0:
            v /= 100*dt
            d += v
            v = np.zeros(3)
    return d/(n//100)

# def calculate_lyapunov(x, y, z, dt):
#     beta = 8/3
#     n = len(x)
#     m = 1
#     v = 0
#     d = 0
#     for i in range(n-1):
#         x_dot, y_dot, z_dot = lorenz(x[i], y[i], z[i])
#         j = x_dot
#         m = j * m
#         v += np.log(np.abs(m))
#         m = m / np.abs(m)
#         if i % 100 == 0:
#             v /= 100*dt
#             d += v
#             v = 0
#     return d/(n//100)

def chua(x, y, z, a = 15.6, b = 1.0, c = 25.58, d = -1, e = 0.02):
    g_x = e * x + (d + e) * (abs(x + 1) - abs(x - 1))
    x_dot = a * (y - x - g_x)
    y_dot = b * (x - y + z)
    z_dot = -c * y
    return x_dot, y_dot, z_dot

def run_chua(x0, y0, z0, t, dt):
    x = np.empty((t,))
    y = np.empty((t,))
    z = np.empty((t,))
    x[0] = x0
    y[0] = y0
    z[0] = z0
    
    for i in range(t-1):
        x_dot, y_dot, z_dot = chua(x[i], y[i], z[i])
        x[i+1] = x[i] + x_dot*dt
        y[i+1] = y[i] + y_dot*dt 
        z[i+1] = z[i] + z_dot*dt
        
    return x, y, z

# 求取Lyapunov指数  
def lyapunov_exponent_lorenz(x0, y0, z0, t, dt, ds0 = 1e-3): 
    # trajectory 1
    x = np.empty((t,))
    y = np.empty((t,))
    z = np.empty((t,))
    x[0] = x0
    y[0] = y0
    z[0] = z0
    for i in range(t-1):
        x_dot, y_dot, z_dot = lorenz(x[i], y[i], z[i])
        x[i+1] = x[i] + x_dot*dt
        y[i+1] = y[i] + y_dot*dt 
        z[i+1] = z[i] + z_dot*dt
    # trajectory 2
    x1 = np.empty((t,))
    y1 = np.empty((t,))
    z1 = np.empty((t,))
    x1[0] = x0 + ds0
    y1[0] = y0
    z1[0] = z0
    for i in range(t-1):
        x_dot, y_dot, z_dot = lorenz(x1[i], y1[i], z1[i])
        x1[i+1] = x1[i] + x_dot*dt
        y1[i+1] = y1[i] + y_dot*dt 
        z1[i+1] = z1[i] + z_dot*dt
    # compute difference
    ds = np.sqrt((x - x1)**2 + (y - y1)**2 + (z - z1)**2)
    lnr = np.log(ds)
    slope, intercept = np.polyfit(list(range(len(lnr))), lnr, deg=1)
    return lnr, slope

def lyapunov_exponent_chua(x0, y0, z0, t, dt, ds0 = 1e-3):
    # trajectory 1
    x = np.empty((t,))
    y = np.empty((t,))
    z = np.empty((t,))
    x[0] = x0
    y[0] = y0
    z[0] = z0
    for i in range(t-1):
        x_dot, y_dot, z_dot = chua(x[i], y[i], z[i])
        x[i+1] = x[i] + x_dot*dt
        y[i+1] = y[i] + y_dot*dt 
        z[i+1] = z[i] + z_dot*dt
    # trajectory 2
    x1 = np.empty((t,))
    y1 = np.empty((t,))
    z1 = np.empty((t,))
    x1[0] = x0 + ds0
    y1[0] = y0
    z1[0] = z0
    for i in range(t-1):
        x_dot, y_dot, z_dot = chua(x1[i], y1[i], z1[i])
        x1[i+1] = x1[i] + x_dot*dt
        y1[i+1] = y1[i] + y_dot*dt 
        z1[i+1] = z1[i] + z_dot*dt
    # compute difference
    ds = np.sqrt((x - x1)**2 + (y - y1)**2 + (z - z1)**2)
    lnr = np.log(ds)
    slope, intercept = np.polyfit(list(range(len(lnr))), lnr, deg=1)
    return lnr, slope

if __name__ == '__main__': 
    x0, y0, z0 = 1,1,1
    t = 5000
    dt = 0.01
    x, y, z = run_lorenz(x0, y0, z0, t, dt)
    # x, y, z = run_chua(x0, y0, z0, t, dt)

    LE = calculate_lyapunov(x, y, z, dt)
    print("lyapunov exponent: ", LE)
    

    # # # lnr, slope = lyapunov_exponent_lorenz(x0, y0, z0, t, dt)
    # # lnr, slope = lyapunov_exponent_chua(x0, y0, z0, t, dt)
    # # print("lyapunov exponent: ", slope)

    # # # plot lnr
    # # plt.figure()
    # # plt.plot(lnr)
    # # plt.title("lyapunov exponent: " + str(slope))
    # # plt.show()
    
    # # plot x, y, z in 3D
    # fig = plt.figure()
    # ax = plt.axes(projection='3d')
    # ax.plot(x, y, z, 'b-', lw=0.5)
    # ax.set_xlabel('x')
    # ax.set_ylabel('y')
    # ax.set_zlabel('z')

    # # line, = ax.plot([], [], [], 'ro', markersize=10)
    # # line.set_markevery([])
    # # # 动画更新函数 
    # # def animate(i): 
    # #     line.set_data(x[i:i+1], y[i:i+1])
    # #     line.set_3d_properties(z[i:i+1])
    # #     line.set_markevery([0])
    # #     return line,
    # # ani = animation.FuncAnimation(fig, animate, frames=x.shape[0], interval=1, blit=False)

    # plt.show()

    # # do fft on x
    # x_fft = np.fft.fft(x)
    # x_fft = np.abs(x_fft)
    # x_fft = x_fft[:int(x_fft.shape[0]/2)]
    # x_fft = x_fft / np.max(x_fft)
    # plt.figure()
    # plt.plot(x_fft)
    # plt.show()
    # # do fft on y
    # y_fft = np.fft.fft(y)
    # y_fft = np.abs(y_fft)
    # y_fft = y_fft[:int(y_fft.shape[0]/2)]
    # y_fft = y_fft / np.max(y_fft)
    # plt.figure()
    # plt.plot(y_fft)
    # plt.show()
    # # do fft on z
    # z_fft = np.fft.fft(z)
    # z_fft = np.abs(z_fft)
    # z_fft = z_fft[:int(z_fft.shape[0]/2)]
    # z_fft = z_fft / np.max(z_fft)
    # plt.figure()
    # plt.plot(z_fft)
    # plt.show()